from mmdet.apis import init_detector, inference_detector
import mmcv
import numpy as np
from tqdm import tqdm
import os
import time
import torch
import argparse
import cv2


def write_result_to_file(result, output_file):
    with open(output_file, "w") as f:
        for k in sorted(result.keys()):
            lines = ["frameIdx " + str(k).zfill(4) + "\n"]
            count = 0
            # for idx, cls in enumerate(["02", "03", "00", "01"]):
            for idx, cls in enumerate(["00", "01"]):
                for item in result[k][idx]:
                    if item[-1] >= result[k][-1]:
                        item[2] -= item[0]
                        item[3] -= item[1]
                        #                         lines.append("class %s position %s \n" % (cls, " ".join(map(lambda x:str( int( (x+32)/64)), item[:-1]))))
                        lines.append(
                            "class %s position %s score %.3f\n"
                            % (
                                cls,
                                " ".join(map(lambda x: str(int(x + 0.5)), item[:-1])),
                                item[-1],
                            )
                        )
                        count += 1
            lines.insert(1, "objNum %d\n" % count)
            lines.append("\n")
            f.writelines(lines)
    print("Record result successfully!")


def main(ckpt_path, video, output_dir, topN=10, ths=0.35, save=False):
    mmcv.mkdir_or_exist(os.path.abspath(output_dir + "/visualize"))
    ckpt = torch.load(ckpt_path)
    with open("tmp.py", "w") as f:
        f.write(ckpt["meta"]["config"])
    model = init_detector("tmp.py", ckpt_path, device="cuda:0")
    print("Init model successfully!")
    os.remove("tmp.py")
    camera = cv2.VideoCapture(video)
    start_time = time.time()
    results = dict()
    frame_id = 1
    while True:
        ret_val, img = camera.read()
        if ret_val:
            result = inference_detector(model, img)
            save_path = os.path.join(output_dir + "/visualize", f"frame_{frame_id}.jpg")
            #             model.show_result(img, result, score_thr=ths, thickness=2, show=False, out_file=save_path)
            result.append(ths)
            results[frame_id] = result
            frame_id += 1
        else:
            break
    camera.release()
    write_result_to_file(results, output_dir + "/results.txt")


def parse_args():
    parser = argparse.ArgumentParser(
        description="generate txt format results",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--ckpt_path", "-cp", help="the model to inference.", type=str, default=None
    )
    parser.add_argument("--video", "-v", help="video.", type=str, default=None)
    parser.add_argument("--topN", "-tn", help="top N scores.", type=int, default=8)
    parser.add_argument(
        "--output_dir", "-o", help="output directory.", type=str, default="./output"
    )

    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = parse_args()
    main(args.ckpt_path, args.video, args.output_dir, args.topN, ths=0.45, save=True)
